package io.kiki.sba.registry.server.clientInterface.store;

import com.google.common.collect.Sets;
import io.kiki.sba.registry.concurrent.UnThrowableCallable;
import io.kiki.sba.registry.util.ConcurrentUtils;
import io.kiki.sba.registry.util.WakeUpLoopExecuteTask;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Collections;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.BiConsumer;

public abstract class DataIndexer<K, V> {
    private static final Logger logger = LoggerFactory.getLogger("SRV-CONNECT");
    private final IndexerRefresher indexerRefresher = new IndexerRefresher();
    private volatile Term lastTerm = new Term();
    private volatile Map<K, Set<V>> index = new ConcurrentHashMap<>(1024);
    private volatile boolean doubleWrite = false;
    private volatile Map<K, Set<V>> tempIndex = new ConcurrentHashMap<>(1024);

    public DataIndexer(String name) {
        ConcurrentUtils.createDaemonThread(name + "-IndexerRefresher", indexerRefresher).start();
    }

    public <R> R add(K key, V val, UnThrowableCallable<R> unThrowableCallable) {
        Term term = lastTerm;
        term.start.incrementAndGet();
        try {
            if (doubleWrite) {
                insert(tempIndex, key, val);
            }
            insert(index, key, val);
            return unThrowableCallable.call();
        } finally {
            term.done.incrementAndGet();
        }
    }

    private void insert(Map<K, Set<V>> map, K key, V val) {
        map.computeIfAbsent(key, k -> Sets.newConcurrentHashSet()).add(val);
    }

    public Set<V> queryByKey(K key) {
        Set<V> set = index.get(key);
        if (set == null) {
            return Collections.emptySet();
        }
        return new HashSet<>(set);
    }

    public Set<K> getKeys() {
        return new HashSet<>(index.keySet());
    }

    private void refresh() {
        tempIndex = new ConcurrentHashMap<>(this.index.size());
        Term prevTerm = lastTerm;
        doubleWrite = true;
        try {
            lastTerm = new Term();
            long startTime = System.currentTimeMillis();
            boolean timeout = !prevTerm.waitAllDone();
            long waitTime = System.currentTimeMillis();
            if (timeout) {
                logger.error("[IndexBuildTimeout]index refresh timeout span={}ms", waitTime - startTime);
            }
            dataStoreForEach((key, val) -> {
                insert(tempIndex, key, val);
            });
            logger.info("index refresh finished waitSpan={}ms, buildSpan={}ms, indexSize={}", waitTime - startTime, System.currentTimeMillis() - waitTime, tempIndex.size());
            index = tempIndex;
        } finally {
            doubleWrite = false;
        }
    }

    public void triggerRefresh() {
        indexerRefresher.wakeup();
    }

    protected abstract void dataStoreForEach(BiConsumer<K, V> biConsumer);

    private static final class Term {
        private static final long timeoutMillis = 10000;
        private final AtomicLong start = new AtomicLong(0);
        private final AtomicLong done = new AtomicLong(0);

        private boolean waitAllDone() {
            long startTs = System.currentTimeMillis();
            for (; ; ) {
                long doneCount = this.done.get();
                long startCount = this.start.get();
                if (startCount == doneCount) {
                    return true;
                }
                if (System.currentTimeMillis() - startTs > timeoutMillis) {
                    return false;
                }
                ConcurrentUtils.sleepUninterruptibly(1, TimeUnit.MILLISECONDS);
            }
        }
    }

    private final class IndexerRefresher extends WakeUpLoopExecuteTask {

        @Override
        public void _execute_() {
            refresh();
        }

        @Override
        public int getWaitingMillis() {
            return (int) ((ThreadLocalRandom.current().nextDouble() * 30 + 30) * 1000);
        }
    }

}
